// test_bn_kernels.cpp
// g++ -std=c++17 test_bn_kernels.cpp -lOpenCL -o test_bn_kernels
// 运行: ./test_bn_kernels

#include <CL/cl.h>
#include <iostream>
#include <vector>
#include <cmath>
#include <cstring>
#include <cassert>
#include <fstream>



inline void checkErr(cl_int err, const char* name) {
    if(err != CL_SUCCESS) {
        std::cerr << "ERROR: " << name << " (" << err << ")\n";
        exit(1);
    }
}

std::string loadKernelSource(const std::string& filename) {
    std::string kernel_path = KERNELS_DIR "/" + filename;
    std::ifstream file(kernel_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open kernel file: " + kernel_path);
    }
    return std::string((std::istreambuf_iterator<char>(file)), 
                      std::istreambuf_iterator<char>());
}

int main() {
    cl_int err;

    // 1) 平台和设备
    cl_uint numPlatforms;
    err = clGetPlatformIDs(0, nullptr, &numPlatforms);
    checkErr(err, "clGetPlatformIDs");
    if(numPlatforms == 0) {
        std::cerr<<"No OpenCL platforms found\n";
        return 1;
    }
    std::vector<cl_platform_id> platforms(numPlatforms);
    err = clGetPlatformIDs(numPlatforms, platforms.data(), nullptr);
    checkErr(err, "clGetPlatformIDs2");

    cl_platform_id platform = platforms[0];

    cl_uint numDevices;
    err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, 0, nullptr, &numDevices);
    checkErr(err, "clGetDeviceIDs count");
    std::vector<cl_device_id> devices(numDevices);
    err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_ALL, numDevices, devices.data(), nullptr);
    checkErr(err, "clGetDeviceIDs");

    cl_device_id device = devices[0];

    // 打印设备名
    {
        char name[256];
        err = clGetDeviceInfo(device, CL_DEVICE_NAME, sizeof(name), name, nullptr);
        if(err==CL_SUCCESS) std::cout<<"Using device: "<<name<<"\n";
    }

    // 2) context & queue
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    checkErr(err, "clCreateContext");
    cl_command_queue q = clCreateCommandQueue(context, device, 0, &err);
    checkErr(err, "clCreateCommandQueue");

    // 3) program
    std::string src = loadKernelSource("bn_utils.cl");
    const char* src_c = src.c_str();
    size_t src_len = src.size();
    cl_program program = clCreateProgramWithSource(context, 1, &src_c, &src_len, &err);
    checkErr(err, "clCreateProgramWithSource");
    err = clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);
    if(err != CL_SUCCESS) {
        // 打印 build log
        size_t log_size=0;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &log_size);
        std::vector<char> log(log_size+1);
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, log.data(), nullptr);
        std::cerr<<"Build log:\n"<<log.data()<<"\n";
        checkErr(err, "clBuildProgram");
    }

    // --- 辅助 lambda: 创建 buffer 并写入 host 数据 ---
    auto createAndWrite = [&](const std::vector<float>& host)->cl_mem {
        cl_mem buf = clCreateBuffer(context, CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR,
                                    sizeof(float)*host.size(), (void*)host.data(), &err);
        checkErr(err, "clCreateBuffer");
        return buf;
    };

    // 现在准备一些测试数据
    const int N = 8;
    // 为通道/批次/空间测试设置小的 dims
    const int batches = 2;
    const int channels = 3;
    const int HW = 4; // 每通道 4 个元素 -> total = batches*channels*HW

    // 简单输入（按索引方便核算）
    std::vector<float> zeros(N, 0.0f);
    std::vector<float> ones(N, 1.0f);
    std::vector<float> mean(N), var(N), gamma(N), beta(N), run_mean(N), run_var(N), a(N), b(N);
    for(int i=0;i<N;++i){
        mean[i] = float(i) * 0.5f;            // 0,0.5,1.0,...
        var[i]  = float(i) * 0.25f + 0.1f;   // avoid zero var
        gamma[i]= 1.0f + 0.1f*float(i);
        beta[i] = 0.2f*float(i);
        run_mean[i] = 0.5f;
        run_var[i]  = 0.2f;
        a[i] = b[i] = 0;
    }

    // 为 forward/backward 使用的 X / Y / dx / dy
    size_t total_elems = size_t(batches) * channels * HW;
    std::vector<float> X(total_elems), Y(total_elems), DX(total_elems, 0.0f), DY(total_elems);
    for(size_t i=0;i<total_elems;++i) {
        X[i] = float(i)*0.1f + 0.5f;
        DY[i] = 0.5f; // 常量梯度
    }

    // buffers for 1D kernels (size N)
    cl_mem buf_mean = createAndWrite(mean);
    cl_mem buf_var  = createAndWrite(var);
    cl_mem buf_gamma= createAndWrite(gamma);
    cl_mem buf_beta = createAndWrite(beta);
    cl_mem buf_run_mean = createAndWrite(run_mean);
    cl_mem buf_run_var  = createAndWrite(run_var);
    cl_mem buf_a = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_a");
    cl_mem buf_b = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_b");

    // buffers for forward/backward
    cl_mem buf_X = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*X.size(), X.data(), &err); checkErr(err,"buf_X");
    cl_mem buf_Y = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*Y.size(), nullptr, &err); checkErr(err,"buf_Y");
    cl_mem buf_DX = clCreateBuffer(context, CL_MEM_READ_WRITE | CL_MEM_COPY_HOST_PTR, sizeof(float)*DX.size(), DX.data(), &err); checkErr(err,"buf_DX");
    cl_mem buf_DY = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*DY.size(), DY.data(), &err); checkErr(err,"buf_DY");

    // per-channel parameters arrays (size = channels)
    std::vector<float> A_ch(channels), B_ch(channels), fx(channels), fdy(channels), b_ch(channels);
    for(int f=0; f<channels; ++f) {
        A_ch[f] = 1.0f + 0.1f*f;
        B_ch[f] = 0.05f * f;
        fx[f] = 0.2f + 0.01f*f;
        fdy[f] = 0.3f + 0.01f*f;
        b_ch[f] = 0.01f * f;
    }
    cl_mem buf_A = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*channels, A_ch.data(), &err); checkErr(err,"buf_A");
    cl_mem buf_B = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*channels, B_ch.data(), &err); checkErr(err,"buf_B");
    cl_mem buf_fx = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*channels, fx.data(), &err); checkErr(err,"buf_fx");
    cl_mem buf_fdy= clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*channels, fdy.data(), &err); checkErr(err,"buf_fdy");
    cl_mem buf_bch= clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*channels, b_ch.data(), &err); checkErr(err,"buf_bch");

    // buffers for compute_backward_factors outputs
    cl_mem buf_x_factor = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_x_factor");
    cl_mem buf_dy_factor = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_dy_factor");
    cl_mem buf_offset = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_offset");

    // buffers for dy_sum / dyx_sum (size N)
    std::vector<float> dy_sum(N, 0.1f), dyx_sum(N, 0.05f);
    cl_mem buf_dy_sum = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*N, dy_sum.data(), &err); checkErr(err,"buf_dy_sum");
    cl_mem buf_dyx_sum= clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float)*N, dyx_sum.data(), &err); checkErr(err,"buf_dyx_sum");

    // buffers for backward_filter outputs
    cl_mem buf_dgamma = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_dgamma");
    cl_mem buf_dbeta  = clCreateBuffer(context, CL_MEM_READ_WRITE, sizeof(float)*N, nullptr, &err); checkErr(err,"buf_dbeta");

    // Helper to run a 1D kernel by name with simple args
    auto run1D = [&](const char* name, cl_kernel kernel, size_t global) {
        err = clEnqueueNDRangeKernel(q, kernel, 1, nullptr, &global, nullptr, 0, nullptr, nullptr);
        checkErr(err, "clEnqueueNDRangeKernel run1D");
        clFinish(q);
    };

    // === 调用 kernels ===

    // 1) update_sums (N)
    {
        cl_kernel k = clCreateKernel(program, "update_sums", &err); checkErr(err,"create update_sums");
        err  = clSetKernelArg(k, 0, sizeof(int), &N);
        err |= clSetKernelArg(k, 1, sizeof(cl_mem), &buf_mean);
        cl_ulong zero = 0;
        err |= clSetKernelArg(k, 2, sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k, 3, sizeof(cl_mem), &buf_var);
        err |= clSetKernelArg(k, 4, sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k, 5, sizeof(cl_mem), &buf_run_mean);
        err |= clSetKernelArg(k, 6, sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k, 7, sizeof(cl_mem), &buf_run_var);
        err |= clSetKernelArg(k, 8, sizeof(cl_ulong), &zero);
        float cmf = 0.9f, rmf = 0.1f, cvf = 0.8f, rvf = 0.2f;
        err |= clSetKernelArg(k, 9, sizeof(float), &cmf);
        err |= clSetKernelArg(k,10, sizeof(float), &rmf);
        err |= clSetKernelArg(k,11, sizeof(float), &cvf);
        err |= clSetKernelArg(k,12, sizeof(float), &rvf);
        checkErr(err,"setargs update_sums");
        size_t g = N;
        run1D("update_sums", k, g);
        std::vector<float> out_run_mean(N), out_run_var(N);
        clEnqueueReadBuffer(q, buf_run_mean, CL_TRUE, 0, sizeof(float)*N, out_run_mean.data(), 0, nullptr, nullptr);
        clEnqueueReadBuffer(q, buf_run_var, CL_TRUE, 0, sizeof(float)*N, out_run_var.data(), 0, nullptr, nullptr);
        std::cout<<"update_sums -> run_mean: ";
        for(float v: out_run_mean) std::cout<<v<<" ";
        std::cout<<"\nupdate_sums -> run_var: ";
        for(float v: out_run_var) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 2) var_gamma_to_a
    {
        cl_kernel k = clCreateKernel(program, "var_gamma_to_a", &err); checkErr(err,"create var_gamma_to_a");
        float eps = 1e-5f;
        cl_ulong zero = 0;
        err  = clSetKernelArg(k,0,sizeof(int), &N);
        err |= clSetKernelArg(k,1,sizeof(float), &eps);
        err |= clSetKernelArg(k,2,sizeof(cl_mem), &buf_var);
        err |= clSetKernelArg(k,3,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,4,sizeof(cl_mem), &buf_gamma);
        err |= clSetKernelArg(k,5,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,6,sizeof(cl_mem), &buf_a);
        err |= clSetKernelArg(k,7,sizeof(cl_ulong), &zero);
        checkErr(err,"setargs var_gamma_to_a");
        size_t g = N;
        run1D("var_gamma_to_a", k, g);
        std::vector<float> out_a(N);
        clEnqueueReadBuffer(q, buf_a, CL_TRUE, 0, sizeof(float)*N, out_a.data(), 0, nullptr, nullptr);
        std::cout<<"var_gamma_to_a -> a: ";
        for(float v: out_a) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 3) mean_var_to_a_b
    {
        cl_kernel k = clCreateKernel(program, "mean_var_to_a_b", &err); checkErr(err,"create mean_var_to_a_b");
        float eps = 1e-5f;
        cl_ulong zero = 0;
        err  = clSetKernelArg(k,0,sizeof(int), &N);
        err |= clSetKernelArg(k,1,sizeof(float), &eps);
        err |= clSetKernelArg(k,2,sizeof(cl_mem), &buf_mean);
        err |= clSetKernelArg(k,3,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,4,sizeof(cl_mem), &buf_var);
        err |= clSetKernelArg(k,5,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,6,sizeof(cl_mem), &buf_a);
        err |= clSetKernelArg(k,7,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,8,sizeof(cl_mem), &buf_b);
        err |= clSetKernelArg(k,9,sizeof(cl_ulong), &zero);
        checkErr(err,"setargs mean_var_to_a_b");
        size_t g = N;
        run1D("mean_var_to_a_b", k, g);
        std::vector<float> out_a(N), out_b(N);
        clEnqueueReadBuffer(q, buf_a, CL_TRUE, 0, sizeof(float)*N, out_a.data(), 0, nullptr, nullptr);
        clEnqueueReadBuffer(q, buf_b, CL_TRUE, 0, sizeof(float)*N, out_b.data(), 0, nullptr, nullptr);
        std::cout<<"mean_var_to_a_b -> a: ";
        for(float v: out_a) std::cout<<v<<" ";
        std::cout<<"\nmean_var_to_a_b -> b: ";
        for(float v: out_b) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 4) combine_mean_var_with_gamma_beta
    {
        cl_kernel k = clCreateKernel(program, "combine_mean_var_with_gamma_beta", &err); checkErr(err,"create combine_mean_var_with_gamma_beta");
        float eps = 1e-5f;
        cl_ulong zero = 0;
        err  = clSetKernelArg(k,0,sizeof(int), &N);
        err |= clSetKernelArg(k,1,sizeof(float), &eps);
        err |= clSetKernelArg(k,2,sizeof(cl_mem), &buf_mean);
        err |= clSetKernelArg(k,3,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,4,sizeof(cl_mem), &buf_var);
        err |= clSetKernelArg(k,5,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,6,sizeof(cl_mem), &buf_gamma);
        err |= clSetKernelArg(k,7,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,8,sizeof(cl_mem), &buf_beta);
        err |= clSetKernelArg(k,9,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,10,sizeof(cl_mem), &buf_a);
        err |= clSetKernelArg(k,11,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,12,sizeof(cl_mem), &buf_b);
        err |= clSetKernelArg(k,13,sizeof(cl_ulong), &zero);
        checkErr(err,"setargs combine_mean_var_with_gamma_beta");
        size_t g = N;
        run1D("combine_mean_var_with_gamma_beta", k, g);
        std::vector<float> out_a(N), out_b(N);
        clEnqueueReadBuffer(q, buf_a, CL_TRUE, 0, sizeof(float)*N, out_a.data(), 0, nullptr, nullptr);
        clEnqueueReadBuffer(q, buf_b, CL_TRUE, 0, sizeof(float)*N, out_b.data(), 0, nullptr, nullptr);
        std::cout<<"combine_mean_var_with_gamma_beta -> a: ";
        for(float v: out_a) std::cout<<v<<" ";
        std::cout<<"\ncombine_mean_var_with_gamma_beta -> b: ";
        for(float v: out_b) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 5) compute_backward_factors
    {
        cl_kernel k = clCreateKernel(program, "compute_backward_factors", &err); checkErr(err,"create compute_backward_factors");
        int M = 10;
        float eps = 1e-5f;
        cl_ulong zero = 0;
        err  = clSetKernelArg(k,0,sizeof(int), &N);
        err |= clSetKernelArg(k,1,sizeof(int), &M);
        err |= clSetKernelArg(k,2,sizeof(float), &eps);
        err |= clSetKernelArg(k,3,sizeof(cl_mem), &buf_mean);
        err |= clSetKernelArg(k,4,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,5,sizeof(cl_mem), &buf_var);
        err |= clSetKernelArg(k,6,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,7,sizeof(cl_mem), &buf_dy_sum);
        err |= clSetKernelArg(k,8,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,9,sizeof(cl_mem), &buf_dyx_sum);
        err |= clSetKernelArg(k,10,sizeof(cl_ulong), &zero);
        // gamma_in: pass gamma
        err |= clSetKernelArg(k,11,sizeof(cl_mem), &buf_gamma);
        err |= clSetKernelArg(k,12,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,13,sizeof(cl_mem), &buf_x_factor);
        err |= clSetKernelArg(k,14,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,15,sizeof(cl_mem), &buf_dy_factor);
        err |= clSetKernelArg(k,16,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,17,sizeof(cl_mem), &buf_offset);
        err |= clSetKernelArg(k,18,sizeof(cl_ulong), &zero);
        checkErr(err,"setargs compute_backward_factors");
        size_t g = N;
        run1D("compute_backward_factors", k, g);
        std::vector<float> out_x(N), out_dy(N), out_of(N);
        clEnqueueReadBuffer(q, buf_x_factor, CL_TRUE, 0, sizeof(float)*N, out_x.data(), 0, nullptr, nullptr);
        clEnqueueReadBuffer(q, buf_dy_factor, CL_TRUE, 0, sizeof(float)*N, out_dy.data(), 0, nullptr, nullptr);
        clEnqueueReadBuffer(q, buf_offset, CL_TRUE, 0, sizeof(float)*N, out_of.data(), 0, nullptr, nullptr);
        std::cout<<"compute_backward_factors -> x_factor: ";
        for(float v: out_x) std::cout<<v<<" ";
        std::cout<<"\ncompute_backward_factors -> dy_factor: ";
        for(float v: out_dy) std::cout<<v<<" ";
        std::cout<<"\ncompute_backward_factors -> offset: ";
        for(float v: out_of) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 6) forward (3D)
    {
        cl_kernel k = clCreateKernel(program, "forward", &err); checkErr(err,"create forward");
        cl_ulong zero = 0;
        err  = clSetKernelArg(k,0,sizeof(int), &batches);
        err |= clSetKernelArg(k,1,sizeof(int), &channels);
        err |= clSetKernelArg(k,2,sizeof(int), &HW);
        err |= clSetKernelArg(k,3,sizeof(cl_mem), &buf_X);
        err |= clSetKernelArg(k,4,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,5,sizeof(cl_mem), &buf_Y);
        err |= clSetKernelArg(k,6,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,7,sizeof(cl_mem), &buf_A);
        err |= clSetKernelArg(k,8,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,9,sizeof(cl_mem), &buf_B);
        err |= clSetKernelArg(k,10,sizeof(cl_ulong), &zero);
        checkErr(err,"setargs forward");
        size_t global[3] = { (size_t)HW, (size_t)channels, (size_t)batches }; // rc, f, b
        err = clEnqueueNDRangeKernel(q, k, 3, nullptr, global, nullptr, 0, nullptr, nullptr);
        checkErr(err,"enqueue forward");
        clFinish(q);
        std::vector<float> outY(total_elems);
        clEnqueueReadBuffer(q, buf_Y, CL_TRUE, 0, sizeof(float)*outY.size(), outY.data(), 0, nullptr, nullptr);
        std::cout<<"forward -> Y: ";
        for(float v: outY) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 7) backward_test (3D)
    {
        // Clear DX to zero first
        std::vector<float> zeroDX(total_elems, 0.0f);
        clEnqueueWriteBuffer(q, buf_DX, CL_TRUE, 0, sizeof(float)*zeroDX.size(), zeroDX.data(), 0, nullptr, nullptr);

        cl_kernel k = clCreateKernel(program, "backward_test", &err); checkErr(err,"create backward_test");
        cl_ulong zero = 0;
        float factor = 0.0f;
        err  = clSetKernelArg(k,0,sizeof(int), &batches);
        err |= clSetKernelArg(k,1,sizeof(int), &channels);
        err |= clSetKernelArg(k,2,sizeof(int), &HW);
        err |= clSetKernelArg(k,3,sizeof(cl_mem), &buf_DX);
        err |= clSetKernelArg(k,4,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,5,sizeof(cl_mem), &buf_DY);
        err |= clSetKernelArg(k,6,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,7,sizeof(cl_mem), &buf_a);
        err |= clSetKernelArg(k,8,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,9,sizeof(float), &factor);
        checkErr(err,"setargs backward_test");
        size_t global[3] = { (size_t)HW, (size_t)channels, (size_t)batches };
        err = clEnqueueNDRangeKernel(q, k, 3, nullptr, global, nullptr, 0, nullptr, nullptr);
        checkErr(err,"enqueue backward_test");
        clFinish(q);
        std::vector<float> outDX(total_elems);
        clEnqueueReadBuffer(q, buf_DX, CL_TRUE, 0, sizeof(float)*outDX.size(), outDX.data(), 0, nullptr, nullptr);
        std::cout<<"backward_test -> DX: ";
        for(float v: outDX) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 8) backward_data (3D)
    {
        // clear DX
        std::vector<float> zeroDX(total_elems, 0.0f);
        clEnqueueWriteBuffer(q, buf_DX, CL_TRUE, 0, sizeof(float)*zeroDX.size(), zeroDX.data(), 0, nullptr, nullptr);

        cl_kernel k = clCreateKernel(program, "backward_data", &err); checkErr(err,"create backward_data");
        cl_ulong zero = 0;
        float factor = 0.0f;
        err  = clSetKernelArg(k,0,sizeof(int), &batches);
        err |= clSetKernelArg(k,1,sizeof(int), &channels);
        err |= clSetKernelArg(k,2,sizeof(int), &HW);
        err |= clSetKernelArg(k,3,sizeof(cl_mem), &buf_X);
        err |= clSetKernelArg(k,4,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,5,sizeof(cl_mem), &buf_DY);
        err |= clSetKernelArg(k,6,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,7,sizeof(cl_mem), &buf_fx);
        err |= clSetKernelArg(k,8,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,9,sizeof(cl_mem), &buf_fdy);
        err |= clSetKernelArg(k,10,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,11,sizeof(cl_mem), &buf_bch);
        err |= clSetKernelArg(k,12,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,13,sizeof(cl_mem), &buf_DX);
        err |= clSetKernelArg(k,14,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,15,sizeof(float), &factor);
        checkErr(err,"setargs backward_data");
        size_t global[3] = { (size_t)HW, (size_t)channels, (size_t)batches };
        err = clEnqueueNDRangeKernel(q, k, 3, nullptr, global, nullptr, 0, nullptr, nullptr);
        checkErr(err,"enqueue backward_data");
        clFinish(q);
        std::vector<float> outDX(total_elems);
        clEnqueueReadBuffer(q, buf_DX, CL_TRUE, 0, sizeof(float)*outDX.size(), outDX.data(), 0, nullptr, nullptr);
        std::cout<<"backward_data -> DX: ";
        for(float v: outDX) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // 9) backward_filter (1D)
    {
        cl_kernel k = clCreateKernel(program, "backward_filter", &err); checkErr(err,"create backward_filter");
        cl_ulong zero = 0;
        float eps = 1e-5f;
        float factor_gamma = 0.0f;
        float factor_beta = 0.0f;
        err  = clSetKernelArg(k,0,sizeof(int), &N);
        err |= clSetKernelArg(k,1,sizeof(cl_mem), &buf_mean);
        err |= clSetKernelArg(k,2,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,3,sizeof(cl_mem), &buf_var);
        err |= clSetKernelArg(k,4,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,5,sizeof(cl_mem), &buf_dy_sum);
        err |= clSetKernelArg(k,6,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,7,sizeof(cl_mem), &buf_dyx_sum);
        err |= clSetKernelArg(k,8,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,9,sizeof(cl_mem), &buf_dgamma);
        err |= clSetKernelArg(k,10,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,11,sizeof(cl_mem), &buf_dbeta);
        err |= clSetKernelArg(k,12,sizeof(cl_ulong), &zero);
        err |= clSetKernelArg(k,13,sizeof(float), &eps);
        err |= clSetKernelArg(k,14,sizeof(float), &factor_gamma);
        err |= clSetKernelArg(k,15,sizeof(float), &factor_beta);
        checkErr(err,"setargs backward_filter");
        size_t g = N;
        run1D("backward_filter", k, g);
        std::vector<float> out_dgamma(N), out_dbeta(N);
        clEnqueueReadBuffer(q, buf_dgamma, CL_TRUE, 0, sizeof(float)*N, out_dgamma.data(), 0, nullptr, nullptr);
        clEnqueueReadBuffer(q, buf_dbeta, CL_TRUE, 0, sizeof(float)*N, out_dbeta.data(), 0, nullptr, nullptr);
        std::cout<<"backward_filter -> dgamma: ";
        for(float v: out_dgamma) std::cout<<v<<" ";
        std::cout<<"\nbackward_filter -> dbeta: ";
        for(float v: out_dbeta) std::cout<<v<<" ";
        std::cout<<"\n";
        clReleaseKernel(k);
    }

    // cleanup
    clReleaseMemObject(buf_mean);
    clReleaseMemObject(buf_var);
    clReleaseMemObject(buf_gamma);
    clReleaseMemObject(buf_beta);
    clReleaseMemObject(buf_run_mean);
    clReleaseMemObject(buf_run_var);
    clReleaseMemObject(buf_a);
    clReleaseMemObject(buf_b);
    clReleaseMemObject(buf_X);
    clReleaseMemObject(buf_Y);
    clReleaseMemObject(buf_DX);
    clReleaseMemObject(buf_DY);
    clReleaseMemObject(buf_A);
    clReleaseMemObject(buf_B);
    clReleaseMemObject(buf_fx);
    clReleaseMemObject(buf_fdy);
    clReleaseMemObject(buf_bch);
    clReleaseMemObject(buf_x_factor);
    clReleaseMemObject(buf_dy_factor);
    clReleaseMemObject(buf_offset);
    clReleaseMemObject(buf_dy_sum);
    clReleaseMemObject(buf_dyx_sum);
    clReleaseMemObject(buf_dgamma);
    clReleaseMemObject(buf_dbeta);

    clReleaseProgram(program);
    clReleaseCommandQueue(q);
    clReleaseContext(context);

    std::cout<<"Done all kernels.\n";
    return 0;
}
